import numpy as np
import matplotlib.pyplot as plt
from TDMA import run_tdma_simulation
from env import ENVIRONMENT

def my_plot(file1, file2, file3):
    # 设置最大迭代次数
    max_iter = 10000

    # 加载已经平滑处理过的奖励数据
    sa_reward_smoothed = np.loadtxt(file1)
    tdma_reward_smoothed = np.loadtxt(file2)
    total_reward_smoothed = np.loadtxt(file3)

    # 绘制吞吐量曲线
    total_line, = plt.plot(np.arange(1, max_iter + 1),total_reward_smoothed,color='#236B8E', lw=1.2,label='Total Throughput')
    sa_line, = plt.plot(np.arange(1, max_iter + 1), sa_reward_smoothed, color='orange', lw=0.8, label='Slotted ALOHA')
    tdma_line, = plt.plot(np.arange(1, max_iter + 1), tdma_reward_smoothed, color='green', lw=0.8, label='TDMA')


    # 设置图例、坐标轴和标题
    plt.legend(handles=[total_line, sa_line, tdma_line])
    plt.xlim((0, max_iter))
    plt.ylim((0, 1.05))
    plt.xlabel('Time Slot')
    plt.ylabel('Throughput')
    plt.title('Network Throughput Comparison')

    # 显示图例
    plt.legend()



my_plot('rewards/only_SA.txt',
        'rewards/only_TDMA.txt',
        'rewards/total.txt'
        )

    # 显示图形
plt.show()
